import os
import cv2

from PIL import ImageGrab
import numpy as np
import time
def read_win():
    last_time = time.time()
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    # video = cv2.VideoWriter('gta_1.mp4',fourcc,20.0,(800,600))
    while True:
        win = cv2.cvtColor(np.array(ImageGrab.grab(bbox=(0,26,800,626))),cv2.COLOR_BGR2RGB)
        print('loop took {} s'.format(time.time()-last_time))
        last_time = time.time()
        cv2.imshow('win',win)
        # video.write(win)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            # video.release()
            cv2.destroyAllWindows()
            break   
def cope_video():
    data_dir = 'data/'
    coped_dir = 'cope_data/'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    for videoname in os.listdir(data_dir):
        read_path = data_dir+videoname
        # print(read_path)
        video_reader = cv2.VideoCapture(read_path)
        ret,frame = video_reader.read()
        output_path = coped_dir+'1080_1920_{}'.format(videoname)
        video_writer = cv2.VideoWriter(output_path,fourcc,20.0,(1920,1080))
        i = 0
        while ret:
            new_frame = frame[:450,...]
            new_frame = cv2.resize(new_frame,(1920,1080))
            # print(new_frame.shape)
            if i%4==0:
                video_writer.write(new_frame)
            ret,frame = video_reader.read()
            i+=1
            # break
        video_writer.release()
        video_reader.release()
        # break

    # print(1080/450,1920/800)
def slow_down():
    data_dir = 'pred_data/'
    video_names = ['1080_1920_gta.mp4','1080_1920_gta_1.mp4']
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    for video_name in video_names:
        video_reader = cv2.VideoCapture(data_dir+video_name)
        video_writer = cv2.VideoWriter(data_dir+'slow_{}'.format(video_name),fourcc,5.0,(1280,720))
        ret,frame = video_reader.read()
        print(ret)
        while ret:
            # print(frame.shape)
            video_writer.write(frame)
            # cv2.imshow('pic',frame)
            # cv2.waitKey(20)
            ret, frame = video_reader.read()
        video_writer.release()
        video_reader.release()
        cv2.destroyAllWindows()
def segment():
    data_dir = 'pred_data/'
    videos= ['slow_1080_1920_gta.mp4','slow_1080_1920_gta_1.mp4']

    # inwindow roi
    pts1 = np.array([[4 * 1280 // 9, 3 * 720 // 4], [5 * 1280 // 9, 3 * 720 // 4], [3 * 1280 // 5, 4 * 720 // 5],
                     [2 * 1280 // 5, 4 * 720 // 5]])
    pts2 = np.asarray([[x - 1280 // 6, y] for (x, y) in pts1])
    pts3 = np.asarray([[x + 1280 // 6, y] for (x, y) in pts1])
    pts4 = np.asarray([[x, y - 720 // 15] for (x, y) in pts1])

    for videoname in videos:
        video = cv2.VideoCapture(data_dir+videoname)
        ret,frame = video.read()
        while ret:
            h,w,_ = frame.shape

            # print(pts2)
            cv2.polylines(frame,[pts1],True,(255,128,0),4)
            cv2.polylines(frame,[pts2],True,(0,0,255),4)
            cv2.polylines(frame,[pts3],True,(255,0,255),4)
            cv2.polylines(frame,[pts4],True,(0,0,0),4)
            mask = np.zeros(frame.shape[:2], dtype=np.uint8)
            # print(mask.shape,frame.shape)
            cv2.fillPoly(mask, np.array([pts1], dtype=np.int32), (255,) * frame.shape[2])

            res = cv2.bitwise_and(frame, frame, mask=mask)
            # cv2.circle(res,(4 * 1280 // 9+10, 3 * 720 // 4+20),5,(255,255,0),-1)
            cv2.imshow('pic',frame[::2,::2])
            cv2.imshow('res',res[::2,::2])
            cv2.waitKey(200)
            print(res.shape,res[3 * 720 // 4+20,4 * 1280 // 9+10 ])
            # break
            ret,frame = video.read()
        video.release()
        break


if __name__=='__main__':
    # read_win()
    # cope_video()
    # slow_down()
    segment()